1770E - Koxia and Tree - CodeForces Solution

#include <bits/stdc++.h>
#define maxn 302333
#define maxm 606666
using namespace std;
using i64 = long long;
const i64 mod = 998244353;
const i64 inv2 = 499122177;
i64 Pow(i64 x, i64 p) {
	i64 r = 1;
	for (; p; p >>= 1, x = x * x % mod) if (p & 1) r = r * x % mod;
	return r;
}
int head[maxn], nxt[maxm], to[maxm], ec = 1, fa[maxn];
int siz[maxn], is[maxn];
void add(int u, int v) {
	nxt[++ec] = head[u], head[u] = ec, to[ec] = v;
}
i64 n, k;
i64 ans = 0;
i64 poss[maxn];
void dfs(int x) {
	siz[x] = is[x];
	for (int e = head[x]; e; e = nxt[e]) {
		if (to[e] != fa[x]) {
			fa[to[e]] = x;
			dfs(to[e]);
			siz[x] += siz[to[e]];
		}
	}
}
void calc_ans(int x) {
	for (int e = head[x]; e; e = nxt[e]) {
		if (to[e] != fa[x]) {
			ans += (i64) siz[to[e]] * (k - siz[to[e]]);
			calc_ans(to[e]);
		}
	}
}
int main() {
	scanf("%lld%lld", &n, &k);
	for (int i = 1, x; i <= k; ++i) scanf("%d", &x), is[x] = 1, poss[x] = 1;
	for (int i = 1, u, v; i < n; ++i) {
		scanf("%d%d", &u, &v);
		add(u, v), add(v, u);
	}
	dfs(1);
	calc_ans(1);
	for (int i = 1, e = 2; i < n; ++i, e += 2) {
		int u = to[e ^ 1], v = to[e];
		if (fa[u] == v) swap(u, v); // u is now father for sure
		ans = (ans + (
				((k - siz[v] - 1) - siz[v] + mod) * poss[u] % mod * (mod + 1 - poss[v]) % mod
				+ ((siz[v] - 1) - (k - siz[v]) + mod) * (mod + 1 - poss[u]) % mod * poss[v] % mod
				) * inv2 % mod) % mod;
		poss[u] = poss[v] = (poss[u] + poss[v]) * inv2 % mod;
	}
//	printf("%lld\n", ans);
	printf("%lld\n", ans * 2 * Pow(k * (k - 1) % mod, mod - 2) % mod);
	return 0;
}